import torch.nn as nn
import torch.nn.functional as F
from Loss.LossRegister import DefaultLossRegister as LossRegister

@LossRegister.register("MSE_loss")
class MSE_loss(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, pred, label):
        if type(pred) == tuple:
            pred = pred[0]
        return F.mse_loss(pred,label)

if __name__ == '__main__':
    import torch
    losser = MSE_loss()
    x = torch.randn((3, 1))

    y = torch.randn((3, 1))
    print(losser(x, y))